"""
This is part of Human Robot Dynamics-Control, By Tourists.

Based on "OpenLoong Dynamics Control"
Transfer from mujoco to ISAAC SIM by IsaacLab

"""
import math
import numpy as np
import mujoco


from ..common import DataBus

TEST_XML = r"""
<mujoco model="test">
  <compiler coordinate="local" angle="radian" eulerseq="xyz"/>
  <size nkey="2"/>
  <option timestep="0.002" gravity="0 0 -9.81"/>
  <visual>
    <global fovy="50" />
    <quality shadowsize="51" />
  </visual>
  <worldbody>
    <geom name="myplane" type="plane" size="10 10 1"/>
    <body name="mybox" pos="0 0 0.1">
      <geom name="mybox" type="box" size="0.1 0.1 0.1" mass="0.25"/>
      <freejoint name="myfree"/>
    </body>
    <body>
      <inertial pos="0 0 0" mass="1" diaginertia="1 1 1"/>
      <site pos="0 0 -1" name="mysite" type="sphere"/>
      <joint name="myhinge" type="hinge" axis="0 1 0" damping="1"/>
    </body>
    <body>
      <inertial pos="0 0 0" mass="1" diaginertia="1 1 1"/>
      <joint name="myball" type="ball"/>
    </body>
    <body mocap="true" pos="42 0 42">
      <geom type="sphere" size="0.1"/>
    </body>
  </worldbody>
  <actuator>
    <position name="myactuator" joint="myhinge"/>
  </actuator>
  <sensor>
    <jointvel name="myjointvel" joint="myhinge"/>
    <accelerometer name="myaccelerometer" site="mysite"/>
  </sensor>
</mujoco>
"""

TEST_XML_SENSOR = r"""
<mujoco model="test">
  <worldbody>
    <geom name="myplane" type="plane" size="10 10 1"/>
  </worldbody>
  <sensor>
    <user objtype="geom" objname="myplane"
          datatype="real" needstage="vel" dim="1"/>
  </sensor>
</mujoco>
"""

TEST_XML_PLUGIN = r"""
<mujoco>
  <option gravity="0 0 0"/>
  <extension>
    <plugin plugin="mujoco.elasticity.cable"/>
  </extension>
  <worldbody>
    <composite type="cable" curve="s" count="41 1 1" size="1" offset="0 0 1" initial="none">
      <plugin plugin="mujoco.elasticity.cable">
        <config key="twist" value="1e6"/>
        <config key="bend" value="1e9"/>
      </plugin>
      <joint kind="main" damping="2"/>
      <geom type="capsule" size=".005" density="1"/>
    </composite>
  </worldbody>
</mujoco>
"""

#include <mujoco/mujoco.h>
#include "data_bus.h"
#include <string>
#include <vector>

class MJ_Interface:
    def __init__(self, mj_modelIn, mj_dataIn):
        # self.jointNum = 0
        self.motor_pos = np.empty(0, dtype=np.float64)
        self.motor_pos_Old = np.empty(0, dtype=np.float64)
        self.motor_vel = np.empty(0, dtype=np.float64)
        self.rpy = np.zeros(3, dtype=np.float64) ## roll,pitch and yaw of baselink
        self.baseQuat = np.zeros(4, dtype=np.float64) ## in quat, mujoco order is [w,x,y,z], here we rearrange to [x,y,z,w]
        self.f3d = np.zeros((3,2), dtype=np.float64) ## 3D foot-end contact force, L for 1st col, R for 2nd col
        self.basePos = np.zeros(3, dtype=np.float64) ##  position of baselink, in world frame
        self.baseAcc = np.zeros(3, dtype=np.float64) ##  acceleration of baselink, in body frame
        self.baseAngVel = np.zeros(3, dtype=np.float64) ##  angular velocity of baselink, in body frame
        self.baseLinVel = np.zeros(3, dtype=np.float64) ##  linear velocity of baselink, in body frame
        self.JointName = [
            "J_arm_l_01","J_arm_l_02","J_arm_l_03", "J_arm_l_04", "J_arm_l_05",
            "J_arm_l_06","J_arm_l_07","J_arm_r_01", "J_arm_r_02", "J_arm_r_03",
            "J_arm_r_04","J_arm_r_05","J_arm_r_06", "J_arm_r_07",
            "J_head_yaw","J_head_pitch","J_waist_pitch","J_waist_roll", "J_waist_yaw",
            "J_hip_l_roll", "J_hip_l_yaw", "J_hip_l_pitch", "J_knee_l_pitch",
            "J_ankle_l_pitch", "J_ankle_l_roll", "J_hip_r_roll", "J_hip_r_yaw",
            "J_hip_r_pitch", "J_knee_r_pitch", "J_ankle_r_pitch", "J_ankle_r_roll"] ## joint name in XML file, the corresponds motors name should be M_*, ref to line 29 of MJ_Interface.cpp
        self.baseName = "base_link"
        self.orientationSensorName = "baselink-quat" ## in quat, mujoco order is [w,x,y,z], here we rearrange to [x,y,z,w]
        self.velSensorName = "baselink-velocity"
        self.gyroSensorName = "baselink-gyro"
        self.accSensorName = "baselink-baseAcc"
        
        # self.mj_model: mujoco.MjModel = mujoco.MjModel.from_xml_string(TEST_XML)
        # self.mj_data = mujoco.MjData(self.mj_model)
        self.mj_model = mj_modelIn
        self.mj_data = mj_dataIn
        # self.jntId_qpos = np.empty(0, dtype=np.int32)
        self.jntId_qvel = np.empty(0, dtype=np.int32)
        self.jntId_dctl = np.empty(0, dtype=np.int32)
        
        self.orientataionSensorId = 0
        self.velSensorId = 0
        self.gyroSensorId = 0
        self.accSensorId = 0
        self.baseBodyId = 0
        
        self.isIni = False
        self.timeStep = 0.001 # second
        self.timeStep = self.mj_model.opt.timestep
        
        self.jointNum = len(self.JointName)
        self.jntId_qpos = np.zeros(self.jointNum)
        self.jntId_qvel = np.zeros(self.jointNum)
        self.jntId_dctl = np.zeros(self.jointNum)
        self.motor_pos = np.zeros(self.jointNum)
        self.motor_vel = np.zeros(self.jointNum)
        self.motor_pos_Old = np.zeros(self.jointNum)

        # MJ_Interface(mjModel *mj_modelIn, mjData  *mj_dataIn);
        # void updateSensorValues();
        # void setMotorsTorque(std::vector<double> &tauIn);
        # void dataBusWrite(DataBus &busIn);
        
        for i in range(self.jointNum):
            tmpId= mujoco.mj_name2id(self.mj_model, mujoco.mjtObj.mjOBJ_JOINT, self.JointName[i])
            if tmpId == -1:
                assert print(self.JointName[i], "  not found in the XML file!")
            self.jntId_qpos[i] = self.mj_model.jnt_qposadr[tmpId]
            self.jntId_qvel[i] = self.mj_model.jnt_dofadr[tmpId]
            motorName = self.JointName[i]
            motorName = "M" + motorName[1:]
            tmpId= mujoco.mj_name2id(self.mj_model, mujoco.mjtObj.mjOBJ_ACTUATOR, motorName)
            if tmpId == -1:
                assert print(motorName, "  not found in the XML file!")
            self.jntId_dctl[i] = tmpId
        # int adr = m->sensor_adr[sensorId];
        # int dim = m->sensor_dim[sensorId];
        # mjtNum sensor_data[dim];
        # mju_copy(sensor_data, &d->sensordata[adr], dim);
        self.baseBodyId= mujoco.mj_name2id(self.mj_model, mujoco.mjtObj.mjOBJ_BODY, self.baseName)
        self.orientataionSensorId= mujoco.mj_name2id(self.mj_model, mujoco.mjtObj.mjOBJ_SENSOR, self.orientationSensorName)
        self.velSensorId= mujoco.mj_name2id(self.mj_model, mujoco.mjtObj.mjOBJ_SENSOR, self.velSensorName)
        self.gyroSensorId= mujoco.mj_name2id(self.mj_model, mujoco.mjtObj.mjOBJ_SENSOR, self.gyroSensorName)
        self.accSensorId= mujoco.mj_name2id(self.mj_model, mujoco.mjtObj.mjOBJ_SENSOR, self.accSensorName)


    def updateSensorValues(self):
        for i in range(self.jointNum):
            self.motor_pos_Old[i] = self.motor_pos[i]
            self.motor_pos[i] = self.mj_data.qpos[int(self.jntId_qpos[i])]
            self.motor_vel[i] = self.mj_data.qvel[int(self.jntId_qvel[i])]
            
        for i in range(4):
            self.baseQuat[i] = self.mj_data.sensordata[self.orientataionSensorId+i]
        tmp = self.baseQuat[0]
        self.baseQuat[0] = self.baseQuat[1]
        self.baseQuat[1] = self.baseQuat[2]
        self.baseQuat[2] = self.baseQuat[3]
        self.baseQuat[3] = tmp

        self.rpy[0] = math.atan2(2*(self.baseQuat[3] * self.baseQuat[0] + self.baseQuat[1] * self.baseQuat[2]), 1-2*(self.baseQuat[0] * self.baseQuat[0] + self.baseQuat[1] * self.baseQuat[1]))
        self.rpy[1] = math.asin(2*(self.baseQuat[3] * self.baseQuat[1] - self.baseQuat[0] * self.baseQuat[2]))
        self.rpy[2] = math.atan2(2*(self.baseQuat[3] * self.baseQuat[2] + self.baseQuat[0] * self.baseQuat[1]), 1-2*(self.baseQuat[1]* self.baseQuat[1] + self.baseQuat[2] * self.baseQuat[2]))

        # for i in range(3):
        #     posOld = self.basePos[i]
        #     self.basePos[i] = self.mj_data.xpos[3 * self.baseBodyId + i].copy()
        #     self.baseAcc[i] = self.mj_data.sensordata[self.accSensorId + i]
        #     self.baseAngVel[i] = self.mj_data.sensordata[self.gyroSensorId + i]
        #     self.baseLinVel[i] = (self.basePos[i] - posOld) / (self.mj_model.opt.timestep)
        
        posOld = self.basePos.copy()
        # Update base position
        self.basePos = self.mj_data.xpos[self.baseBodyId].copy()
        # Update acceleration and angular velocity
        self.baseAcc = self.mj_data.sensordata[self.accSensorId:self.accSensorId+3]
        self.baseAngVel = self.mj_data.sensordata[self.gyroSensorId:self.gyroSensorId+3]
        # Calculate linear velocity
        self.baseLinVel = (self.basePos - posOld) / self.mj_model.opt.timestep

    def setMotorsTorque(self, tauIn):
        for i in range(self.jointNum):
            self.mj_data.ctrl[i]=tauIn.at(i)

    def dataBusWrite( self, busIn):
        busIn.motors_pos_cur =  self.motor_pos
        busIn.motors_vel_cur =  self.motor_vel
        busIn.rpy[0] = self.rpy[0]
        busIn.rpy[1] = self.rpy[1]
        busIn.rpy[2] = self.rpy[2]
        busIn.fL[0] = self.f3d[0][0]
        busIn.fL[1] = self.f3d[1][0]
        busIn.fL[2] = self.f3d[2][0]
        busIn.fR[0] = self.f3d[0][1]
        busIn.fR[1] = self.f3d[1][1]
        busIn.fR[2] = self.f3d[2][1]
        busIn.basePos[0] = self.basePos[0]
        busIn.basePos[1] = self.basePos[1]
        busIn.basePos[2] = self.basePos[2]
        busIn.baseLinVel[0] = self.baseLinVel[0]
        busIn.baseLinVel[1] = self.baseLinVel[1]
        busIn.baseLinVel[2] = self.baseLinVel[2]
        busIn.baseAcc[0] = self.baseAcc[0]
        busIn.baseAcc[1] = self.baseAcc[1]
        busIn.baseAcc[2] = self.baseAcc[2]
        busIn.baseAngVel[0] = self.baseAngVel[0]
        busIn.baseAngVel[1] = self.baseAngVel[1]
        busIn.baseAngVel[2] = self.baseAngVel[2]
        busIn.updateQ()












